import os
import random
from PIL import Image, ImageDraw, ImageFont
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import gymnasium as gym
from gymnasium import spaces

from gymnasium.envs.registration import register

raise NotImplementedError("最佳坐标存在问题，需要调整")

# 随机选择的中文字集合（可以根据需求进行扩展）
CHINESE_CHARACTERS = list("汉字随机生成测试常用文字")

# 切换到 'TkAgg' 后端，确保图形显示在本地窗口
matplotlib.use('TkAgg')  # 或者使用 'Agg' 后端避免任何显示窗口


class ImageTextGame(gym.Env):
    def __init__(self, image_dir, font_paths, max_text_count=4, render_mode='human'):
        super(ImageTextGame, self).__init__()
        self.image_dir = image_dir
        self.font_paths = font_paths  # 多个字体文件路径的列表
        self.max_text_count = max_text_count  # 限制最大的文字数量
        self.window_width = 260  # 宽度为260
        self.window_height = 160  # 高度保持不变
        self.grid_width = 20  # 横向20等分
        self.grid_height = 10  # 纵向10等分
        self.current_image = None
        self.text_info = []  # 保存多个文字的相关信息
        self.current_step = 0  # 当前的点击步数
        self.total_reward = 0  # 总得分
        self.max_steps = 0  # 最大步数，等于绘制的文字数
        self.click_positions = []  # 记录玩家的点击位置
        self.render_mode = render_mode  # 控制是否可视化
        self.text_string = ""  # 保存绘制的文字串

        # 定义离散的动作空间，200个方格 (20 x 10)
        self.action_space = spaces.Discrete(self.grid_width * self.grid_height)

        # 定义观察空间，图像的像素值范围在 0 到 255 之间
        self.observation_space = spaces.Box(low=0, high=255, shape=(self.window_height, self.window_width, 3),
                                            dtype=np.uint8)

    def reset(self, seed=None, options=None):
        # 重置游戏状态
        super().reset(seed=seed)
        self.current_image = self._get_random_image()
        self.text_info = []  # 确保在reset时清空之前的文字信息
        self.current_image = self._draw_random_text(self.current_image)
        self.current_step = 0
        self.total_reward = 0
        self.max_steps = len(self.text_info)  # 最大步数等于绘制的文字数量
        if self.max_steps == 0:
            raise ValueError("max_steps cannot be 0, check if text was drawn correctly.")
        self.click_positions = []  # 清空点击位置记录

        # 返回(图片, 文字字符串)
        return (self.current_image, self.text_string), {}

    def step(self, action):
        '''
        这里需要后期考虑是否需要即使用户点击了区域不在文字区域内
        但是如果该区域也占用了部分文字区域，也算点击成功
        参考链接：https://chatgpt.com/c/66f8ee7d-9858-800b-b7c2-578425e34391
        '''
        if isinstance(action, list):
            action = action[0]
        # 将离散动作映射到网格位置
        grid_x = action % self.grid_width
        grid_y = action // self.grid_width
        # print(f"grid_x: {grid_x}, grid_y: {grid_y}, action:{action}")

        # 计算网格单元的尺寸
        grid_cell_width = self.window_width / self.grid_width
        grid_cell_height = self.window_height / self.grid_height

        # 计算点击位置为网格单元的中心
        x_click = int(grid_x * grid_cell_width)
        y_click = int(grid_y * grid_cell_height)

        # 检查游戏是否已终止
        if self.current_step >= self.max_steps:
            terminated = True
            truncated = False
            return (self.current_image, self.text_string), self.total_reward, terminated, truncated, {}

        target_info = self.text_info[self.current_step]

        # 获取旋转字符的位置信息和尺寸
        x_pos, y_pos = target_info['position']
        text_width = target_info['width']
        text_height = target_info['height']

        current_reward = 0
        # 检查点击是否在旋转字符的边界框内
        if x_pos <= x_click <= x_pos + text_width and y_pos <= y_click <= y_pos + text_height:
            # 点击正确
            self.total_reward += 2
            current_reward = 2
        else:
            # 点击错误
            self.total_reward -= 2
            current_reward = -2

        # 记录点击位置
        self.click_positions.append((x_click, y_click))

        # 在图片上标记所有点击位置（不包括方框，只标记红点）
        img = Image.fromarray(self.current_image.copy())
        draw = ImageDraw.Draw(img)
        for (click_x, click_y) in self.click_positions:
            draw.ellipse([click_x - 5, click_y - 5, click_x + 5, click_y + 5], fill="red", outline="red")

        # 更新 obs 中的图片，包含点击位置
        updated_image = np.array(img)

        # 增加步骤
        self.current_step += 1

        # 检查终止条件
        terminated = self.current_step >= self.max_steps
        truncated = False

        if self.render_mode == 'human':
            self.render()

        return (updated_image, self.text_string), current_reward, terminated, truncated, {}


    def render(self, mode='human'):
        img = Image.fromarray(self.current_image.copy())
        draw = ImageDraw.Draw(img)

        # 标记当前目标字符的边界框
        if self.current_step < len(self.text_info):
            target_info = self.text_info[self.current_step]
            x_pos, y_pos = target_info['position']
            text_width = target_info['width']
            text_height = target_info['height']
            draw.rectangle([x_pos, y_pos, x_pos + text_width, y_pos + text_height], outline="green", width=2)

        # 标记所有点击位置
        for (x_click, y_click) in self.click_positions:
            draw.ellipse([x_click - 5, y_click - 5, x_click + 5, y_click + 5], fill="red", outline="red")

        # 显示图像
        plt.imshow(np.array(img))
        plt.title(f"Step: {self.current_step}, Total Reward: {self.total_reward}")
        plt.axis('off')
        plt.show(block=False)
        plt.pause(0.001)


    def _get_random_image(self):
        # 随机选择目录中的一张图片
        image_files = os.listdir(self.image_dir)
        selected_image_path = os.path.join(self.image_dir, random.choice(image_files))
        image = Image.open(selected_image_path)

        # 调整图片大小
        image = self._adjust_image_size(image)
        return np.array(image)

    def _adjust_image_size(self, image):
        # 获取图片的尺寸
        img_width, img_height = image.size

        if img_width < self.window_width or img_height < self.window_height:
            # 如果图片的大小小于窗口大小，则拉伸图片
            image = image.resize((self.window_width, self.window_height))
        else:
            # 如果图片的大小大于窗口大小，随机截取一个区域
            x_start = random.randint(0, img_width - self.window_width)
            y_start = random.randint(0, img_height - self.window_height)
            image = image.crop((x_start, y_start, x_start + self.window_width, y_start + self.window_height))

        return image

    def _draw_random_text(self, image):
        # 将numpy数组转换为PIL图像
        img_pil = Image.fromarray(image)
        draw = ImageDraw.Draw(img_pil)
        self.text_string = ""  # 重置文本字符串

        # 随机生成四个汉字
        text = random.choices(CHINESE_CHARACTERS, k=4)

        x_position = 0

        for char in text:
            self.text_string += char

            font_size = random.randint(int(self.window_height * 0.15), int(self.window_height * 0.35))
            font_path = random.choice(self.font_paths)
            font = ImageFont.truetype(font_path, font_size)

            color = tuple(np.random.randint(0, 255, size=3))

            # 创建字符图像
            bbox = draw.textbbox((0, 0), char, font=font)
            text_width = bbox[2] - bbox[0]
            text_height = bbox[3] - bbox[1]

            char_img = Image.new("RGBA", (text_width, text_height), (255, 255, 255, 0))
            char_draw = ImageDraw.Draw(char_img)
            char_draw.text((0, 0), char, font=font, fill=color)

            # 旋转字符图像
            angle = random.randint(-45, 45)
            rotated_char_img = char_img.rotate(angle, expand=1)

            # 获取非透明像素的边界框
            rotated_bbox = rotated_char_img.getbbox()
            # 将旋转字符图像裁剪到边界框
            cropped_rotated_char_img = rotated_char_img.crop(rotated_bbox)

            # 更新旋转后的宽度和高度
            rotated_width, rotated_height = cropped_rotated_char_img.size

            # 调整 x_position 和 y_position 以确保字符在图像边界内
            if x_position + rotated_width > self.window_width:
                x_position = self.window_width - rotated_width
            if x_position < 0:
                x_position = 0

            y_position = random.randint(0, self.window_height - rotated_height)
            if y_position + rotated_height > self.window_height:
                y_position = self.window_height - rotated_height
            if y_position < 0:
                y_position = 0

            # 将裁剪后的旋转字符图像粘贴到主图像上
            img_pil.paste(cropped_rotated_char_img, (x_position, y_position), cropped_rotated_char_img)

            # 记录字符信息
            self.text_info.append({
                "char": char,
                "font_size": font_size,
                "color": color,
                "position": (x_position, y_position),
                "font_path": font_path,
                "rotation_angle": angle,
                "width": rotated_width,
                "height": rotated_height,
            })

            x_position += rotated_width + 10

        return np.array(img_pil)

    
    def get_target_area(self):
        if self.current_step < len(self.text_info):
            target_info = self.text_info[self.current_step]

            x_pos, y_pos = target_info['position']
            text_width = target_info['width']
            text_height = target_info['height']

            grid_cell_width = self.window_width / self.grid_width
            grid_cell_height = self.window_height / self.grid_height

            x_pos_min = x_pos
            x_pos_max = x_pos + text_width - 1
            y_pos_min = y_pos
            y_pos_max = y_pos + text_height - 1

            grid_x_min = max(0, int(np.ceil(x_pos_min / grid_cell_width)))
            grid_x_max = min(self.grid_width - 1, int(np.floor(x_pos_max / grid_cell_width)))
            grid_y_min = max(0, int(np.ceil(y_pos_min / grid_cell_height)))
            grid_y_max = min(self.grid_height - 1, int(np.floor(y_pos_max / grid_cell_height)))

            # 确保 grid_x_max >= grid_x_min 且 grid_y_max >= grid_y_min
            if grid_x_max < grid_x_min:
                grid_x_max = grid_x_min
            if grid_y_max < grid_y_min:
                grid_y_max = grid_y_min

            grid_x = random.randint(grid_x_min, grid_x_max)
            grid_y = random.randint(grid_y_min, grid_y_max)

            action = grid_y * self.grid_width + grid_x

            # print(f"grid_x: {grid_x}, grid_y: {grid_y}, action:{action}")

            return (self.current_image, self.text_string), action
        else:
            return None, None


def save_obs(obs, size = None):
    # 保存obs中的图片，图片名字以obs中的文字串命名
    image, text_string = obs
    img = Image.fromarray(image)

    if size:
        img = img.resize(size, Image.LANCZOS)

    filename = f"{text_string}.png"
    img.save(filename)


def test_get_target_area():
    env = ImageTextGame(image_dir=r"D:\Projects\datasets\gym\wordVision", font_paths=[r"D:\Projects\datasets\gym\font\simsun.ttc"], render_mode="human")
    for _ in range(1000):
        obs, _ = env.reset()
        # 使用标准的 gym 动作生成方式
        for i in range(env.max_steps):
            _, action = env.get_target_area()  # 随机生成点击位置
            obs, reward, _, _, _ = env.step(action)
            # 保存obs中的图片，图片名字以obs中的文字串命名
            save_obs(obs)
            print(f"Current reward: {reward}")
            if reward < 0:
                raise ValueError("Game Over")

    print("Test Pass")


def test_env():
    env = ImageTextGame(image_dir=r"D:\Projects\datasets\gym\wordVision", font_paths=[r"D:\Projects\datasets\gym\font\simsun.ttc"], render_mode="human")
    for _ in range(1):
        obs, _ = env.reset()
        # 使用标准的 gym 动作生成方式
        for i in range(env.max_steps):
            action = env.action_space.sample()  # 随机生成点击位置
            obs, reward, terminated, truncated, _ = env.step(action)
            save_obs(obs, size=(84, 84))
            print(f"Current reward: {reward}")
            if terminated or truncated:
                print("Game Over")
                break

    print("Text Info:", env.text_info)  # 输出绘制的文字信息


# 用法示例
if __name__ == "__main__":
    # 注册自定义环境，id是唯一标识
    register(
        id='ImageTextGame-v0',  # 环境ID
        entry_point='game.ImageTextGame',  # 指向你环境的入口，即类的导入路径
        max_episode_steps=4,  # 每一回合最大步数
        kwargs={'image_dir': r"D:\Projects\datasets\gym\wordVision",  # 初始化时传入的参数
                'font_paths': [r"D:\Projects\datasets\gym\font\simsun.ttc"],
                'max_text_count': 4,
                'render_mode': 'service'}
    )

    # test_env()
    test_get_target_area()
